#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @author: Gaoxiang

import torch
import os
import logging
import numpy as np
import pandas as pd
from text_cls_config import Config
from text_cls_dataloader import DataIterator
from text_cls_dataloader import dataloader_init
from text_cls_model import fastText

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# raw_data = pd.read_csv("./data_test - 副本.csv")
# # raw_data = pd.read_csv("./data_test.csv")
# raw_data = raw_data.rename(columns={"flag": "flag_col", "上行短信内容": "content_col"})
raw_data = "./data_test.csv"

def main(config, raw_data):
    if not os.path.isdir(config["model_path"]):
        os.mkdir(config["model_path"])

    config["to_build_new_vocab"] = False
    datafield = dataloader_init(config=config)
    predict_data = DataIterator(raw_set=raw_data,
                                is_test_set=True,
                                use_cuda=config["use_cuda"],
                                batchsize=config["batch_size"],
                                datafield=datafield,
                                config=config)

    model = fastText(config)
    if config["use_cuda"]:
        model = model.cuda()
    model.load_state_dict(torch.load("epoch_100 (3).pth", map_location=torch.device('cpu')))
    # model.load_state_dict(torch.load(config["model_path"] + "/epoch_99.pth"))
    model.eval()

    for data in predict_data.get_tier_data():
        if config["use_cuda"]:
            content = data.text.cuda()
        else:
            content = data.text
        outputs = model(content)
        outputs = torch.nn.functional.softmax(outputs, dim=1)
        prob, predictions = torch.max(outputs.data, 1)
        predictions_np = predictions.detach().cpu().numpy()
        prob_np = prob.detach().cpu().numpy()
        np_mapping = np.vectorize(lambda x: config["label_map"][x])
        result_np = np_mapping(predictions_np)
        print(result_np, prob_np)
    print(type(result_np))


if __name__ == "__main__":
    main(Config, raw_data)
